package com.ihai.spring.boot.sample.zookeeper;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.TreeCache;
import org.apache.curator.framework.recipes.cache.TreeCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.utils.CloseableUtils;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.UriSpec;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.io.Closeable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

/**
 * zookeeper service
 *
 * @author 张驰
 * @date 2019/6/10
 */
public class ZkService {
    private static final Logger logger = LoggerFactory.getLogger(ZkService.class);
    private CuratorFramework client;
    private String path;
    private ServiceDiscovery<String> serviceDiscovery;
    private List<TreeCache> treeCacheList = new ArrayList<>();

    public void start(String host, String port, String path, String serviceName, String url) throws Exception {
        this.path = path;
        client = CuratorFrameworkFactory.newClient(host + ":" + port, 1000, 30000, new ExponentialBackoffRetry(1000, 3));
        client.start();
        if (StringUtils.hasText(url)) {
            serviceDiscovery = ServiceDiscoveryBuilder.builder(String.class).client(client).basePath(path).build();
        } else {
            registServiceInstance(path, serviceName, url);
        }
        serviceDiscovery.start();
    }

    public void start(String host, String port, String path) throws Exception {
        start(host, port, path, null, null);
    }

    public List<ServiceInstance<String>> getInstance(String serviceName) throws Exception {
        try {
            Collection<ServiceInstance<String>> instances = serviceDiscovery.queryForInstances(serviceName);
            return (List<ServiceInstance<String>>) CollectionUtils.arrayToList(instances.toArray(new ServiceInstance[0]));
        } finally {
            CloseableUtils.closeQuietly(serviceDiscovery);
        }
    }

    public void watch(String serviceName, TreeCacheListener treeCacheListener) throws Exception {
        TreeCache treeCache = new TreeCache(client, path + "/" + serviceName);
        treeCache.start();
        treeCache.getListenable().addListener(treeCacheListener);
        treeCacheList.add(treeCache);
    }

    private void registServiceInstance(String path, String serviceName, String url) throws Exception {
        UriSpec uriSpec = new UriSpec(url);
        ServiceInstance<String> thisInstance = ServiceInstance.<String>builder()
                .name(serviceName)
                .payload(url)
                .uriSpec(uriSpec)
                .build();

        serviceDiscovery = ServiceDiscoveryBuilder.builder(String.class)
                .client(client)
                .basePath(path)
                .thisInstance(thisInstance)
                .build();
    }

    public void close() {
        for (TreeCache treeCache : treeCacheList) {
            close(treeCache);
        }
        close(serviceDiscovery);
        close(client);
    }

    private void close(Closeable closeable) {
        try {
            if (closeable != null) {
                closeable.close();
            }
        } catch (Exception e) {
            logger.warn(closeable + " closed with exception. this can be ignored.", e);
        }
    }
}
